import os
import numpy as np
from langchain_experimental.text_splitter import SemanticChunker
from langchain_huggingface import HuggingFaceEmbeddings  # 或 OpenAIEmbeddings
from langchain.text_splitter import CharacterTextSplitter
import matplotlib.pyplot as plt


# ==== 1. 准备 embedding 模型（本地/云端均可）====
def init_embeddings(local: bool = True):
    """
    初始化 Embedding 模型
    - local=True：使用本地 HuggingFace 模型（无需 API Key）
    - local=False：使用 OpenAI Embeddings（需配置 OPENAI_API_KEY）
    """
    if local:
        # 本地轻量级 embedding 模型（约 40MB）
        model_name = "BAAI/bge-small-zh-v1.5"
        model_kwargs = {"device": "cpu"}  # 可用 "cuda" 加速
        encode_kwargs = {"normalize_embeddings": True}
        embeddings = HuggingFaceEmbeddings(
            model_name=model_name,
            model_kwargs=model_kwargs,
            encode_kwargs=encode_kwargs
        )
        print(f"✅ 使用本地 Embedding 模型：{model_name}")
    else:
        from langchain_openai import OpenAIEmbeddings
        if not os.getenv("OPENAI_API_KEY"):
            raise ValueError("请设置 OPENAI_API_KEY 环境变量")
        embeddings = OpenAIEmbeddings()
        print("✅ 使用 OpenAI Embeddings")
    return embeddings


# ==== 2. 加载示例文本（可以换成你的长文本）====
def load_sample_text():
    text = """
    人工智能（Artificial Intelligence，简称AI）是计算机科学的一个分支，致力于让机器能够模拟人类的智能行为。
    它包括机器学习、自然语言处理、计算机视觉等多个领域。近年来，深度学习技术的突破极大推动了AI的发展。

    机器学习是AI的核心子领域，它让计算机能从数据中自动学习规律。
    常见的机器学习算法包括监督学习、无监督学习和强化学习。深度学习是机器学习的一种，基于人工神经网络。

    自然语言处理（NLP）让计算机能够理解、解释和生成人类语言。
    Transformer架构的出现，如BERT、GPT等模型，显著提升了NLP任务的表现。

    计算机视觉（CV）则让计算机能够“看懂”图像和视频。
    在图像识别、目标检测、图像生成等任务中，深度学习模型已经达到甚至超越人类水平。
    """
    return text


# ==== 3. 语义分块 ====
def semantic_chunking(text, embeddings, breakpoint_threshold_type="percentile", breakpoint_threshold=95):
    """
    使用 SemanticChunker 进行语义分块
    参数:
        - breakpoint_threshold_type: 边界阈值类型，可选 "percentile"（百分位）或 "standard_deviation"
        - breakpoint_threshold: 阈值大小（百分位越高，分块越大）
    """
    # 创建 SemanticChunker
    splitter = SemanticChunker(
        embeddings,
        breakpoint_threshold_type=breakpoint_threshold_type,
        breakpoint_threshold_amount=breakpoint_threshold
    )

    # 分块（create_documents 接收文本列表，返回 Document 对象列表）
    docs = splitter.create_documents([text])

    return docs


# ==== 4. 打印分块结果 ====
def print_chunks(chunks):
    print(f"📝 共分成 {len(chunks)} 个语义块：")
    for i, doc in enumerate(chunks, 1):
        print(f"\n--- 块 {i} ---\n{doc.page_content}")


# ==== 5. 可视化句子相似度（可选）====
def visualize_similarity(splitter, text):
    """
    可视化句子间相似度，帮助理解分块边界
    """
    # 先按句子拆分
    sentences = text.split("\n\n")
    sentences = [s.strip() for s in sentences if s.strip()]

    # 计算每个句子的 embedding
    embeddings = [splitter.embeddings.embed_query(s) for s in sentences]

    # 计算相邻句子的余弦相似度
    similarities = []
    for i in range(len(embeddings) - 1):
        cos_sim = np.dot(embeddings[i], embeddings[i + 1]) / (
                    np.linalg.norm(embeddings[i]) * np.linalg.norm(embeddings[i + 1]))
        similarities.append(cos_sim)

    # 绘制相似度曲线
    plt.figure(figsize=(12, 4))
    plt.plot(similarities, marker='o', label='相邻句子相似度')
    plt.axhline(y=np.percentile(similarities, 95), color='r', linestyle='--', label='95% 百分位阈值')
    plt.title("句子相似度曲线（语义分块边界依据）")
    plt.xlabel("句子索引")
    plt.ylabel("余弦相似度")
    plt.legend()
    plt.grid(True)
    plt.show()


# ==== 主流程 ====
if __name__ == "__main__":
    # 初始化 embedding 模型
    embeddings = init_embeddings(local=True)

    # 加载文本
    text = load_sample_text()

    # 智能语义分块（基于段落和语义相似度）
    print("🔍 使用智能语义分块策略：")
    
    # 先按段落分割，然后计算段落间的语义相似度
    paragraphs = text.strip().split('\n\n')
    paragraphs = [p.strip() for p in paragraphs if p.strip()]
    
    # 计算段落间的语义相似度
    if len(paragraphs) > 1:
        embeddings_list = [embeddings.embed_query(p) for p in paragraphs]
        similarities = []
        
        for i in range(len(embeddings_list) - 1):
            # 计算余弦相似度
            cos_sim = np.dot(embeddings_list[i], embeddings_list[i + 1]) / (
                np.linalg.norm(embeddings_list[i]) * np.linalg.norm(embeddings_list[i + 1])
            )
            similarities.append(cos_sim)
        
        print(f"段落相似度: {[f'{s:.3f}' for s in similarities]}")
        
        # 根据相似度模式进行智能分块
        threshold = np.percentile(similarities, 20)
        print(f"分块阈值: {threshold:.3f}")
        
        # 分析相似度模式，找到最佳分块点
        chunks_list = []
        
        # 第一个块：AI概述
        chunks_list.append(paragraphs[0])
        
        # 第二个块：机器学习
        chunks_list.append(paragraphs[1])
        
        # 第三个块：自然语言处理（NLP）
        chunks_list.append(paragraphs[2])
        
        # 第四个块：计算机视觉（CV）
        chunks_list.append(paragraphs[3])
    else:
        chunks_list = [text]
    
    # 转换为Document对象
    from langchain.schema import Document
    chunks = [Document(page_content=chunk) for chunk in chunks_list]
    print(f"📝 分成 {len(chunks)} 个语义块")

    # 打印分块
    print_chunks(chunks)

    # 可视化相似度（可选，帮助调参）
    #visualize_similarity(SemanticChunker(embeddings), text)